# 1. A proxy model is like a post processor that can run on your data before it can reach the views. It can do processing like
# 'sorting' and 'filtering'. Processing that requires sorting or filtering should use the QSortFilterProxyModel class.
# If you need a special behavior that is neither sorting nor filtering that you should look into subclassing QAbstractProxyModel
# to provide your own processing methods.
# 2. Proxy model is always attached between the view and the model:  Data --> Model --> Proxy --> View.
# 3. ProxyModel doesn't change the data of the structure. It just remaps the QModelIndex-es and sends them off to the view.
# In this example we will use the QSortFilterProxyModel.
# 4. There are two ways to load a Qt Designer made UI. a) By generating an intermediate python file and loading it.
# b) Loading the inheritable classes at runtime. We will do the latter. For that we need to import the pysideuic (for PySide)
# or the uic (for Qt4). UIC stands for User Interface Compiler and it originally belongs to the PyQt4 package. 
# Also note that PySide does not have the loadUiType method that is why we have used a user implementation of that method.

import sys
from PySide import QtCore, QtGui
import icons_rc                 # Import the icons file that was made earlier
from loadUiType import *

class Node(object):
    def __init__(self, name, parent = None):
        self._name = name       # Name of the node
        self._children = []     # List of Children of the node
        self._parent = parent   # Parent of the node
        
        # If there is a parent to this node, then this node must add itself to the child list of the parent node
        if parent is not None:
            parent.addChild(self)
        
    # typeInfo method will return a string which will describe the type of this item. Since it is a node, it will return type as Node
    def typeInfo(self):
        return "Node"
    
    # Add a new child to the list of children
    def addChild(self, child):
        self._children.append(child)    # Append the new child to the list of children using append method.
        
    # Similar to the add child method above but puts the child at a given arbitrary location.
    def insertChild(self, position, child):
        # We check if the position is within the boundary of the list. If not, we return false.
        if position < 0 or position > len(self._children):
            return False
        # Else put the child at the given position using the insert method. Set the parent of the child to self and return true 
        self._children.insert(position, child)
        child._parent = self
        return True
    
    # Just like the insert child method, the remove child method is implemented as follows.
    def removeChild(self, position):
        # We check if the position is within the boundary of the list. If not, we return false.
        if position < 0 or position > len(self._children):
            return False
        # Else get the child at the given position using the pop method. Then et the parent of the child to None and return true 
        child = self._children.pop(position)
        child._parent = None
        return True    

    # Method to return the name of an object
    def name(self):
        return self._name
    
    # Method to set the node name as required by the setData method
    def setName(self, name):
        self._name = name
    
    # Method to get the child corresponding to an index (or row)
    def child(self, row):
        return self._children[row]
    
    # Method to get the number of children
    def childCount(self):
        return len(self._children)
    
    # Method to get the parent of this node
    def parent(self):
        return self._parent
    
    # Method to get the index number of this node relative to its parent
    def row(self):
        if self._parent is not None:
            return self._parent._children.index(self)
        
    # Method to display the data in the console
    def log(self, tabLevel = -1):
        output = ""
        tabLevel += 1
        
        for i in range(tabLevel):
            output += "\t"
        
        output += self._name + "\n"
        
        for child in self._children:
            output += child.log(tabLevel)
        
        tabLevel -= 1
        return output
    
    def __repr__(self):
        return self.log()
    
    
# Let's make some subclasses from the Node class and call them LightNode, TransformNode, CameraNode. All of them would need to 
# reimplement the typeInfo method. Initially we forced each class to have a parent. But now we changed it be without a parent.
class LightNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(LightNode, self).__init__(name, parent)
        
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "LightNode"

class TransformNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(TransformNode, self).__init__(name, parent)
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "TransformNode"

class CameraNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(CameraNode, self).__init__(name, parent)
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "CameraNode"


# Now we will implement the tree model. To implement the tree model, we have to subclass QAbstractItemModel and provide methods for it.
class SceneGraphModel(QtCore.QAbstractItemModel):
    """ Inputs: Node, QObject """
    def __init__(self, root, parent = None):
        super(SceneGraphModel, self).__init__(parent)
        self._rootNode = root
    
    """ Inputs: QModelIndex """
    """ Outputs: int """
    """ Return the amount of children an item has """
    """ Note that we don't need the parent argument in the 'list' and the 'table' model. Only for 'tree' models. """
    """ We get the rootNode if the parent is not valid, otherwise we get the parent. Then we just return the childCount of the node """
    def rowCount(self, parent):
        if not parent.isValid():
            parentNode = self._rootNode
        else:
            parentNode = parent.internalPointer()
        
        return parentNode.childCount()  
    
    """ Inputs: QModelIndex """
    """ Outputs: int """
    """ Column count controls the number of columns in the tree view. Giving a column count of 2 makes two columns in the header of the tree view """
    """ But it does not display the other sections in tree view style, that's because the parent method of our model returns 0 on the second """
    """ parameter, which is actually the column of the createIndex method. This is telling that make column 0 (the first column) a hierarchical """ 
    """ view only. Now how should we use the second column. We can use the second column to display the type of the node. """ 
    """ Sounds like a good idea. For that we have to change the data and the headerData methods """
    def columnCount(self, parent):
        return 2
        
    # Data method to display something on the view. We will return both text and icons through this method.
    """ Inputs: QModelIndex, int """
    """ Outputs: QVariant, strings are cast to QString which is a QVariant """
    # Remember that the data method returns something that the view wants to display. All we will do here is call interalPointer to get our node,
    # then if it is displayrole, return the node name.
    # If the index is valid, we pass the test. Get the node. If it's displayrole, return the node name.
    # The index is a QModelIndex and it stores information about where an item is. Column corresponds to the header section.
    def data(self, index, role):
        if not index.isValid():
            return None
        
        node = index.internalPointer()
        
        # What to do when the view wants to display data, is done thorough DisplayRole.
        # Similarly, what to do in case when the view wants to edit data is done through the EditRole.
        # * Another thing to Note here: we haven't explicitly set the roles for sorting and filtering, so it defaults to using
        # the displayRole. Since we are returning the node name while in displayRole, node name will become the basis for
        # sorting and filtering.
        if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole: # to retain value when item is double clicked for editing
            if index.column() == 0:     # if item lies in zeroth column then return the node name
                return node.name()
            if index.column() == 1:     # if item is in first column, return the node 'typeinfo'
                return node.typeInfo()  # to display the node typeInfo in the second column
        
        # Icons are returned by using a decoration role.
        if role == QtCore.Qt.DecorationRole:
            if index.column() == 0:     # if item lies in zeroth column then first get the 'typeinfo'. Depending on that, return the icon.
                typeInfo = node.typeInfo()
                if typeInfo == "LightNode":
                    # For returning an icon, we have to construct one using the QtGui.QIcon() method. The QIcon takes a QPixmap as an input.
                    # So we construct that as well. The QPixmap, takes a resource as input. To retrieve a resource at runtime, we type
                    # :/resource_name. Note that we do not want to work with full paths here.
                    return QtGui.QIcon(QtGui.QPixmap(":/Light.png"))
                if typeInfo == "CameraNode":
                    return QtGui.QIcon(QtGui.QPixmap(":/Camera.png"))
                if typeInfo == "TransformNode":
                    return QtGui.QIcon(QtGui.QPixmap(":/Transform.png"))
    
    # HeaderData method to display something on the view's header
    """ Inputs: int, Qt::Orientation, int """
    """ Outputs: QVariant, strings are cast to QString which is a QVariant """
    """ Let's just return the string "SceneGraph" to display on the header """
    def headerData(self, section, orientation, role):
        if role == QtCore.Qt.DisplayRole:
            if section == 0:
                return "SceneGraph"
            else:
                return "TypeInfo"   # To show TypeInfo in the header
    
    # Flags method to select the items, or enable items so that they are not greyed out etc. 
    """ Inputs: QModelIndex """
    """ Outputs: int (flag) """
    """ Flags method, as in previous examples, needs to return that the items are enabled, the items are selectable and the items are editable """
    """ The itemIsEditable flag requires that we implement a method called setData """
    def flags(self, index):
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable
    
    """ Inputs: QModelIndex, QVariant, int (flag) """
    """ The setData method receives an index (QModelIndex), a value (QVariant - A QVariant is a data holder for Qt applications) and a role """
    """ We want the role to default to EditRole """
    """ We will use the internalPointer to get our node and then set its name """
#     def setData(self, index, value, role = QtCore.Qt.EditRole):
#         if index.isValid():
#             if role == QtCore.Qt.EditRole:
#                 node = index.internalPointer()
#                 node.setName(value)
#                 return True
#         return False
    def setData(self, index, value, role = QtCore.Qt.EditRole):
        if index.isValid():
            node = index.internalPointer()
            if role == QtCore.Qt.EditRole:
                # If the index column is 0 (i.e. the node name column), then change the name.
                # Also calling setData manually wouldn't update the views that are using this model. They wouldn't notice
                # that the data has changed. So we need to emit a signal, which we get by inheriting QAbstractItemModel.
                # The signal is called dataChanged. It takes two parameters, from and to. Since we only want to update a
                # single item, we will pass the same index to both of them.  
                if index.column() == 0:
                    node.setName(value)
                    self.dataChanged.emit(index, index)
                return True
        return False

    # Tree models and Tree view require two more special methods
    """ Remember that the QModelIndex class has valuable information about where and item was with the row and columns methods that return ints"""
    """ In the previous examples, we used those methods index our lists. """
    """ The QModelIndex also have something called and index pointer. The indexPointer returns a custom node instance which we implemented above """
    # Our goal of from this method is to return the parent of the given index and the returned value must be a QModelIndex class itself
    """ Inputs: QModelIndex """
    """ Outputs: QModelIndex """
    """ Returns the parent of the node with the given QModelIndex """
    def parent(self, index):
        node = self.getNode(index) # Returns a node
        parentNode = node.parent()
        # If the returned node is the 'root' node then return an empty QModelIndex because roots don't have a parent
        if parentNode == self._rootNode:
            return QtCore.QModelIndex()
        # If the returned node is not the root node, then it must be a child node. Now is the time to create a QModelIndex class object that
        # we will wrap the parent node into. This is done by the factor method called 'createIndex'. We have this method because of inheriting
        # from the QAbstractItemModel. First arg is an integer that corresponds to the row that the parent node lies in, so we implement the
        # custom method row() that we created above. The second arg can be 0 for now. The third argument is an object that can be anything. It's
        # also the object that will be returned when you call internal pointer of the index that this method returns. We will pass our parent node
        # to this argument.
        # self.createIndex(int, int, object=0)
        return self.createIndex(parentNode.row(), 0, parentNode)

    """ Inputs: int, int, QModelIndex """
    """ Outputs: QModelIndex """
    """ Returns the QModelIndex that corresponds to the given row, column and parent node """
    # The row and column are integers that the view acquires from the rowCount and the columnCount methods. It also receives a parent which it
    # acquires from the method implemented above. The index method is responsible for returning a child at the given row and given column of the
    # given parent.   
    def index(self, row, column, parent):
        parentNode = self.getNode(parent)
        
        # Now we have the parent node. After that, we ask the parentNode to get the child by calling our own child method and passing the row 
        # that we want the child from. The row is an integer.
        childItem = parentNode.child(row)
        
        # But if for some wierd reason the child is None, we return an empty QModelIndex. If child exists we create an index and return it.
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex

    # Let's make a custom method, that is actually not overriding anything.
    # It will act like a helper. If we pass a QModelIndex with a valid index, we will get the internal pointer (i.e. the node,
    # if it exists) and return it. Otherwise we return the root node.  
    # First check if the parent is valid. Because in the 'parent' method above notice that we return an empty QModelIndex in case of a root
    # node, which is not valid. Therefore we set the parent node here to rootNode if the parent is not valid i.e. an empty QModelIndex. 
    # But if parent is valid then we get the internal pointer which again corresponds to our node instance.
    """ Custom method (not overriding anything) """
    """ Inputs: QModelIndex """
    def getNode(self, index):
        if index.isValid():
            node = index.internalPointer()
            if node:
                return node
        else:
            return self._rootNode
        
    # Methods for Inserting and Removing rows at runtime.
    # Recall that in the last examples, we ignored the parent arguments. Now is the time we will use it.
    # Insert rows acts as an insert children method, i.e. you are basically inserting children to this parent.
    # For that we will make an insert child method in the Node class which will add a child at an arbitrary location. 
    """ Inputs: int, int, QModelIndex """
    def insertRows(self, position, rows, parent = QtCore.QModelIndex()):
        # To insert the children, we need to get the parent node and we get it by calling our customly implemented getNode method
        # by passing the parent QModelIndex class to it. As a result it returns its internalPointer which is the name of the node.
        parentNode = self.getNode(parent)
    # Inputs are parent to insert at, the position to insert at, and how many items (more accurately, the ending index of items)
        self.beginInsertRows(parent, position, position + rows - 1)
        for row in range(rows):
            # Keep track of child count
            childCount = parentNode.childCount()
            # Create a node
            childNode = Node("NewNode" + str(childCount))
            # Insert this child node. The current child node pushes the previous child nodes, downwards 
            success = parentNode.insertChild(position, childNode)
        self.endInsertRows()
        return success
    
    """ Inputs: int, int, QModelIndex """
    def removeRows(self, position, rows, parent = QtCore.QModelIndex()):
        parentNode = self.getNode(parent)
        self.beginRemoveRows(parent, position, position + rows - 1)
        for row in range(rows):
            success = parentNode.removeChild(position)
        self.endRemoveRows()
        return success
    
    # As a demonstration we can also insert light nodes
    def insertLightRows(self, position, rows, parent = QtCore.QModelIndex()):
        parentNode = self.getNode(parent)
        self.beginInsertRows(parent, position, position + rows - 1)
        for row in range(rows):
            # Keep track of child count
            childCount = parentNode.childCount()
            # Create a node
            childNode = LightNode("NewLightNode" + str(childCount))
            # Insert this child node. The current child node pushes the previous child nodes, downwards 
            success = parentNode.insertChild(position, childNode)
        self.endInsertRows()
        return success
    
# The uic module contains a method called loadUiType (here I have used a custom user implementation of that module).
# We input the .ui file saved from the Qt Designer to this function. This method returns the form class and the base class,
# which we need to derive from, and create a new window.
base, form = loadUiType("myUI.ui")

# Let's create a new class that derives from the form and the base class and call that Wnd. It is derived from form and base class.
class Wnd(base, form):
    # In the constructor, we need to call the super class constructor of the base class and pass the parent to it. Then we also
    # need to call the setupUi method that gets generated by the loadUiType method. setupUi method actually initializes all our
    # controls placed in the Qt Designer. 
    def __init__(self, parent=None):
        super(base, self).__init__(parent)
        self.setupUi(self)
        
        rootNode = Node("Head")
        childNode1 = Node("Torso", rootNode)
        childNode2 = Node("RightArm", childNode1)
        childNode3 = CameraNode("LeftArm", childNode1)
        childNode4 = Node("RightWrist", childNode2)
        childNode5 = TransformNode("LeftWrist", childNode3)
        childNode6 = Node("RightHand", childNode2)
        childNode7 = Node("LeftHand", childNode3)
        childNode8 = LightNode("RightElbow", childNode2)
        childNode8 = Node("LeftElbow", childNode3)
        
        # Implement the filter
        # Our goal is to type something in the lineEdit and based on that, the uiTree should show the filtered results.
        # For that we have to use the QSortFilterProxyModel class.  So let's create one.        
        self._proxyModel = QtGui.QSortFilterProxyModel()
        # Now we have our proxy model. The proxy model needs to be in between view and the data model i.e. 
        # View <------> Proxy Model <------> Data Model. So instead of hooking our model directly into the view as below
        # model = SceneGraphModel(rootNode, self)   # <------ This is the data model
        # self.uiTree.setModel(model)     # <------ This is the view. You can see it is directly hooked to each other.
        # We put our proxy model in between. So instead of hooking the model directly into the view, we will hook the model
        # into the proxy model and then take the proxy model and hook it into the view. Also we made the models as protected
        # members using the self._model notations.

        self._model = SceneGraphModel(rootNode, self)   # <------ This is the data model
        self._proxyModel.setSourceModel(self._model) # <------ This is the proxy model
        self._proxyModel.setDynamicSortFilter(True) # To enable dynamic sorting as text is entered
        self._proxyModel.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive) # To enable case insensitivity 
        self.uiTree.setModel(self._proxyModel)     # <------ This is the view
        # All we need to do now is to call the setFilterRegExp with the required string.
        # First let's add some nodes. We want the added nodes in the root to we specify no parent.
        self._model.insertLightRows(0, 10)
        # For filtering out data using hardcoded filter string we can use the following code
        # self._proxyModel.setFilterRegExp("Light")
        # But we want to use our lineEdit's (we named the lineEdit, uiFilter) current text. For that we need to use signals
        # and slots provided by the Qt framework.
        # LineEdit emits the signal textChanged when the text is changed by the user. The signal also holds a parameter with the
        # current text. So if we would connect that signal to the setFilterRegExp method, it will automatically set the new filter
        # text whenever we edit the filter field. To connect the textChanged signal to the setFilterRegExp method we use this code.
        # connect() is a static function. First parameter is the signal emitter, second is.. which signal we want to connect i.e.
        # here we have the textChanged signal with type QString. Third is whose the receiver of that signal. Here it is the proxy
        # model and the function from the proxy model is setFilterRegExp.
        QtCore.QObject.connect(self.uiFilter, QtCore.SIGNAL("textChanged(QString)"), self._proxyModel.setFilterRegExp)
        
        # Now let's implement sorting
        # For that we only need to call setSortingEnabled method on our uiTree.
        # Note: It would show an error if we already don't have a proxy model in between the data model and the view. 
        self.uiTree.setSortingEnabled(True)
        
        # If you need a differnt logic or process than QSortFilterProxyModel then you should consider subclassing the
        # QtGui.QAbstractProxyModel and implement the behavior you are after.        

        
if __name__ == "__main__":
    
    appObject = QtGui.QApplication(sys.argv)
    appObject.setStyle('cleanlooks')
        
    # Create the class and show it
    wnd = Wnd()
    wnd.show()
    
    # Execute the application
    sys.exit(appObject.exec_())
